import os
import cv2
import numpy as np
import glob

img_path = r"I:\0_Deep_Learning\磁瓦\度量学习\001_ciwa_dataset\val\wuzi_04"
save_dir = r"I:\0_Deep_Learning\磁瓦\度量学习\001_ciwa_dataset\pad_val\wuzi_04"
assert img_path != save_dir

# img_list = glob.glob(os.path.join(img_path, "*.jpg"))
img_list = os.listdir(img_path)


def show_img(img):
    cv2.imshow('img', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

i = 0
for p in img_list:
    if p.endswith('.jpg'):
        path = os.path.join(img_path, p)
        img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
        h, w, _ = img.shape
        # target_size = (224, 224)

        if w > h:
            square = np.zeros((w, w, 3), dtype=np.uint8)

            delta = int((w - h) / 2)
            # 图片填充位置坐标
            # crop_x1 = 0
            crop_y1 = delta
            # crop_x2 = w - 1
            crop_y2 = h + delta

            square[crop_y1:crop_y2, :, :] = img
            # show_img(square)
        else:
            square = np.zeros((h, h, 3), dtype=np.uint8)

            delta = int((h - w) / 2)
            crop_x1 = delta
            crop_x2 = w + delta

            square[:, crop_x1:crop_x2, :] = img
            # show_img(square)
        
        save_path = os.path.join(save_dir, p)
        # cv2.imwrite(save_path, square)
        cv2.imencode('.jpg', square)[1].tofile(save_path)
        print(p, "\t填充成功")
        i += 1

print('共 ', i, ' 张')